# 作者:小宇
# 2025年11月12日13时44分06秒
# 1225074067@qq.com
import pandas as pd
import json
import re
from pathlib import Path


def extract_disaster_data_from_text(text):
    """从文本中提取灾害相关信息"""
    if pd.isna(text) or not isinstance(text, str):
        return None

    # 提取灾害年份的模式
    year_patterns = [
        r'(\d{3,4})\s*年',
        r'公元\s*(\d{3,4})',
        r'前\s*(\d+)\s*年',
        r'(\d{3,4})—(\d{3,4})',  # 年份范围
        r'(\d{3,4})至(\d{3,4})',
        r'(\d{3,4})～(\d{3,4})',
    ]

    # 提取灾害影响的模式
    impact_patterns = [
        r'死[伤亡毙]\s*(\d+(?:\.\d+)?)\s*人',
        r'溺毙[死伤]\s*(\d+(?:\.\d+)?)\s*人',
        r'压死[伤]\s*(\d+(?:\.\d+)?)\s*人',
        r'死亡\s*(\d+(?:\.\d+)?)\s*人',
        r'伤\s*(\d+(?:\.\d+)?)\s*人',
        r'溺死\s*(\d+(?:\.\d+)?)\s*人',
        r'(\d+)\s*余人死亡',
        r'死亡\s*(\d+)\s*余人',
    ]

    # 提取房屋损毁的模式
    house_patterns = [
        r'坏[毁损]房[屋舍]\s*(\d+(?:\.\d+)?)\s*间',
        r'塌房[屋舍]\s*(\d+(?:\.\d+)?)\s*间',
        r'房屋倒塌\s*(\d+(?:\.\d+)?)\s*间',
        r'庐舍倾颓\s*(\d+(?:\.\d+)?)\s*间',
        r'冲塌民房\s*(\d+(?:\.\d+)?)\s*间',
    ]

    # 提取农田损毁的模式
    farmland_patterns = [
        r'坏[损毁]田[稼禾]\s*(\d+(?:\.\d+)?)\s*顷',
        r'淹没田[稼禾]\s*(\d+(?:\.\d+)?)\s*顷',
        r'冲淹田[稼禾]\s*(\d+(?:\.\d+)?)\s*顷',
        r'田[稼禾]荡尽',
        r'禾稼尽没',
        r'颗粒无收',
    ]

    years = []
    death_count = None
    house_damage = None
    farmland_damage = None

    # 提取年份
    for pattern in year_patterns:
        matches = re.findall(pattern, text)
        for match in matches:
            if isinstance(match, tuple):  # 年份范围
                try:
                    start_year = int(match[0])
                    end_year = int(match[1])
                    years.extend([start_year, end_year])
                except:
                    pass
            else:  # 单个年份
                try:
                    year = int(match)
                    years.append(year)
                except:
                    pass

    # 提取死亡人数
    for pattern in impact_patterns:
        match = re.search(pattern, text)
        if match:
            try:
                death_count = int(match.group(1))
                break
            except:
                pass

    # 提取房屋损毁数量
    for pattern in house_patterns:
        match = re.search(pattern, text)
        if match:
            try:
                house_damage = int(match.group(1))
                break
            except:
                pass

    # 提取农田损毁面积
    for pattern in farmland_patterns:
        match = re.search(pattern, text)
        if match:
            try:
                farmland_damage = int(match.group(1)) if match.group(1).isdigit() else None
                break
            except:
                pass

    # 去重并排序年份
    years = sorted(list(set(years)))

    return {
        'years': years,
        'death_count': death_count,
        'house_damage': house_damage,
        'farmland_damage': farmland_damage
    }


def extract_disaster_data_from_excel(file_path, sheet_name='总灾害'):
    """从Excel文件中提取灾害数据"""
    try:
        # 读取Excel文件
        df = pd.read_excel(file_path, sheet_name=sheet_name)

        results = []

        for idx, row in df.iterrows():
            # 获取时期（朝代）
            period = row.iloc[0] if not pd.isna(row.iloc[0]) else None

            # 获取灾害类型
            disaster_types = row.iloc[3] if len(row) > 3 else None

            # 获取灾害时间
            disaster_time = row.iloc[4] if len(row) > 4 else None

            # 获取灾害后果
            consequences = row.iloc[5] if len(row) > 5 else None

            # 获取原文描述（用于提取详细信息）
            description = row.iloc[1] if len(row) > 1 else None

            if period and disaster_types:
                # 处理灾害类型（可能包含多个类型，用顿号、逗号分隔）
                if pd.isna(disaster_types):
                    types_list = []
                else:
                    # 分割灾害类型
                    types_list = re.split(r'[、，, ]', str(disaster_types))
                    types_list = [t.strip() for t in types_list if t.strip()]

                # 从原文和后果中提取详细信息
                extracted_info = {}
                if description:
                    extracted_info = extract_disaster_data_from_text(str(description))
                if consequences and not extracted_info.get('death_count'):
                    consequence_info = extract_disaster_data_from_text(str(consequences))
                    if consequence_info.get('death_count'):
                        extracted_info['death_count'] = consequence_info['death_count']

                # 创建灾害记录
                record = {
                    'period': str(period),
                    'disaster_types': types_list,
                    'disaster_time': str(disaster_time) if not pd.isna(disaster_time) else None,
                    'consequences': str(consequences) if not pd.isna(consequences) else None,
                    'years': extracted_info.get('years', []),
                    'death_count': extracted_info.get('death_count'),
                    'house_damage': extracted_info.get('house_damage'),
                    'farmland_damage': extracted_info.get('farmland_damage')
                }

                # 只添加有灾害类型的记录
                if types_list:
                    results.append(record)

        return results

    except Exception as e:
        print(f"读取文件时出错: {e}")
        return []


def main():
    # 输入文件路径
    input_file = './data/05灾害 - 总数据和各朝代数据.xlsx'

    # 输出文件路径
    output_file = './outputs/disaster_data.json'

    print("开始提取灾害数据...")

    # 提取灾害数据
    data = extract_disaster_data_from_excel(input_file, sheet_name='总灾害')

    print(f"从'总灾害' sheet提取了 {len(data)} 条记录")

    # 保存为JSON文件
    Path(output_file).parent.mkdir(parents=True, exist_ok=True)
    with open(output_file, 'w', encoding='utf-8') as f:
        json.dump(data, f, ensure_ascii=False, indent=2)

    print(f"数据已保存到: {output_file}")

    # 统计信息
    total_deaths = sum([r.get('death_count', 0) for r in data if r.get('death_count')])
    total_house_damage = sum([r.get('house_damage', 0) for r in data if r.get('house_damage')])

    print(f"\n统计信息:")
    print(f"  总记录数: {len(data)}")
    print(f"  总死亡人数: {total_deaths}")
    print(f"  总房屋损毁: {total_house_damage} 间")

    # 按灾害类型统计
    disaster_type_count = {}
    for record in data:
        for disaster_type in record['disaster_types']:
            disaster_type_count[disaster_type] = disaster_type_count.get(disaster_type, 0) + 1

    print(f"\n灾害类型分布:")
    for disaster_type, count in sorted(disaster_type_count.items(), key=lambda x: x[1], reverse=True)[:10]:
        print(f"  {disaster_type}: {count} 次")

    # 显示前几条数据作为示例
    print("\n前5条数据示例:")
    for i, record in enumerate(data[:5], 1):
        print(f"\n记录 {i}:")
        print(f"  时期(period): {record['period']}")
        print(f"  灾害类型(disaster_types): {record['disaster_types']}")
        print(f"  灾害时间(disaster_time): {record['disaster_time']}")
        print(f"  死亡人数(death_count): {record['death_count']}")
        print(f"  房屋损毁(house_damage): {record['house_damage']} 间")
        print(f"  农田损毁(farmland_damage): {record['farmland_damage']} 顷")
        print(f"  涉及年份(years): {record['years']}")

    return data


if __name__ == '__main__':
    main()